#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
# 手眼合一色块抓取
import cv2
import numpy as np
import time
import threading
import signal
import LeArm
import kinematics as kin
import RPi.GPIO as GPIO

debug = True
ball_color = 'green'
stream = "http://127.0.0.1:8080/?action=stream?dummy=param.mjpg"
cap = cv2.VideoCapture(stream)

# 逆运动学参数
F = 1000/240.0  # 像素到舵机脉冲的转换系数
center_x = 320   # 图像中心X坐标
center_y = 240   # 图像中心Y坐标

# PID控制器
x_pid = kin.PID(P=0.01, I=0.001, D=0)
y_pid = kin.PID(P=0.00001, I=0, D=0)

# 机械臂控制参数
grasp_height = 200    # 抓取高度（单位：mm）
approach_dist = 50    # 接近距离（单位：mm）
gripper_open = 1200   # 爪子张开位置
gripper_close = 700   # 爪子闭合位置

# 坐标转换参数
cam_offset_x = 0      # 摄像头X方向偏移（机械臂末端到摄像头）
cam_offset_y = 0      # 摄像头Y方向偏移
cam_offset_z = 100    # 摄像头Z方向偏移（单位：mm）

orgFrame = None
Running = False
position_color_list = []
cv_blocks_ok = False

# 校准按键
GPIO.setmode(GPIO.BCM)
GPIO.setwarnings(False)
key = 22
GPIO.setup(key, GPIO.IN, GPIO.PUD_UP)

color_dist = {
    'red': {'Lower': np.array([0, 60, 60]), 'Upper': np.array([6, 255, 255])},
    'blue': {'Lower': np.array([100, 80, 46]), 'Upper': np.array([124, 255, 255])},
    'green': {'Lower': np.array([35, 43, 35]), 'Upper': np.array([90, 255, 255])},
}

def signal_handler(sig, frame):
    global Running
    print("Stopping...")
    Running = False
    cv2.destroyAllWindows()
    LeArm.runActionGroup('rest', 1)
    exit(0)

signal.signal(signal.SIGINT, signal_handler)

def image_processing():
    global orgFrame, position_color_list, cv_blocks_ok
    while Running:
        ret, frame = cap.read()
        if not ret:
            continue
        try:
            frame = cv2.resize(frame, (640, 480))
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            mask = cv2.inRange(hsv, color_dist[ball_color]['Lower'], color_dist[ball_color]['Upper'])
            mask = cv2.erode(mask, None, iterations=2)
            mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_RECT, (3,3)), iterations=2)
            
            contours = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
            if len(contours) > 0:
                c = max(contours, key=cv2.contourArea)
                rect = cv2.minAreaRect(c)
                box = cv2.boxPoints(rect)
                cv2.drawContours(frame, [np.int0(box)], -1, (0,255,0), 2)
                
                (x, y), (w, h), angle = rect
                if w*h < 1000:  # 过滤小面积噪声
                    continue
                
                # 转换到机械臂坐标系
                target_x, target_y = pixel_to_arm(x, y)
                position_color_list = [(target_x, target_y, angle)]
                cv_blocks_ok = True
                
            cv2.imshow('Frame', frame)
            cv2.waitKey(1)
        except Exception as e:
            print("Image processing error:", e)

def pixel_to_arm(px, py):
    """将像素坐标转换为机械臂坐标系"""
    # 获取当前机械臂末端位置
    current_pos = kin.get_current_position()
    
    # 计算相对于图像中心的偏移
    dx = (px - center_x) * 0.1  # 比例系数需根据实际校准
    dy = (py - center_y) * 0.1
    
    # 转换为机械臂坐标系（考虑摄像头安装位置）
    target_x = current_pos[0] + dx + cam_offset_x
    target_y = current_pos[1] + dy + cam_offset_y
    return target_x, target_y

def arm_control():
    global cv_blocks_ok
    LeArm.runActionGroup('rest', 1)
    time.sleep(1)
    
    while Running:
        if cv_blocks_ok and position_color_list:
            x, y, angle = position_color_list[0]
            
            # 计算抓取位姿
            target_pos = (x, y, grasp_height)
            approach_pos = (x, y + approach_dist, grasp_height)
            
            # 逆运动学求解
            servo_data = kin.ik_transform(target_pos)
            if not servo_data:
                print("IK求解失败!")
                continue
                
            # PID控制调整
            x_pid.SetPoint = x
            y_pid.SetPoint = y
            x_pid.update(servo_data['x'])
            y_pid.update(servo_data['y'])
            
            # 执行运动
            try:
                # 接近目标
                kin.move_to(approach_pos)
                time.sleep(1)
                
                # 精确定位
                for _ in range(3):  # 三次调整
                    kin.servo_adjust(x_pid.output, y_pid.output)
                    time.sleep(0.1)
                
                # 抓取动作
                LeArm.setServo(1, gripper_close, 500)
                time.sleep(0.5)
                kin.move_to((x, y, grasp_height + 50))  # 抬升
                time.sleep(1)
                
                # 返回初始位置
                LeArm.runActionGroup('rest', 1)
                
            except Exception as e:
                print("运动控制错误:", e)
            
            cv_blocks_ok = False
            position_color_list = []
            
        time.sleep(0.1)

if __name__ == '__main__':
    Running = True
    LeArm.runActionGroup('rest', 1)
    
    # 启动图像处理线程
    img_thread = threading.Thread(target=image_processing)
    img_thread.daemon = True
    img_thread.start()
    
    # 启动控制线程
    ctrl_thread = threading.Thread(target=arm_control)
    ctrl_thread.daemon = True
    ctrl_thread.start()
    
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        Running = False
        cv2.destroyAllWindows()
        LeArm.runActionGroup('rest', 1)